import pandas as pd
import numpy as np
import os

def add_initial_capacity_to_file(file_path):
    """为电池数据文件添加初始容量列"""
    
    # 读取Excel文件
    df = pd.read_excel(file_path)
    
    print(f"处理文件: {file_path}")
    print(f"原始数据形状: {df.shape}")
    print(f"列名: {df.columns.tolist()}")
    
    # 识别不同的电池（基于循环次数重置）
    # 当循环次数从高值突然变为1时，表示新的电池开始
    battery_ids = []
    current_battery = 0
    prev_cycle = 0
    
    for cycle in df['cycle']:
        if cycle == 1 and prev_cycle > 1:
            current_battery += 1
        battery_ids.append(current_battery)
        prev_cycle = cycle
    
    df['battery_id'] = battery_ids
    
    # 为每个电池计算初始容量（循环次数为1时的容量）
    initial_capacities = []
    
    for idx, row in df.iterrows():
        battery_id = row['battery_id']
        cycle = row['cycle']
        
        # 找到该电池循环次数为1的记录
        battery_data = df[df['battery_id'] == battery_id]
        initial_capacity_row = battery_data[battery_data['cycle'] == 1]
        
        if not initial_capacity_row.empty:
            initial_capacity = initial_capacity_row['Capacity'].iloc[0]
        else:
            # 如果没有找到循环次数为1的记录，使用该电池的第一个循环容量
            first_cycle_row = battery_data[battery_data['cycle'] == battery_data['cycle'].min()]
            initial_capacity = first_cycle_row['Capacity'].iloc[0] if not first_cycle_row.empty else row['Capacity']
        
        initial_capacities.append(initial_capacity)
    
    # 添加初始容量列
    df['initial_capacity'] = initial_capacities
    
    print(f"处理后数据形状: {df.shape}")
    print(f"电池数量: {df['battery_id'].nunique()}")
    print(f"初始容量统计:")
    print(f"  最小值: {df['initial_capacity'].min():.2f}")
    print(f"  最大值: {df['initial_capacity'].max():.2f}")
    print(f"  平均值: {df['initial_capacity'].mean():.2f}")
    
    # 保存处理后的数据
    output_path = file_path.replace('.xlsx', '_with_initial_capacity.xlsx')
    df.to_excel(output_path, index=False)
    print(f"处理后的数据已保存到: {output_path}")
    
    # 显示每个电池的初始容量
    print("\n各电池初始容量:")
    battery_summary = df.groupby('battery_id').agg({
        'initial_capacity': 'first',
        'cycle': 'count'
    }).rename(columns={'cycle': 'cycle_count'})
    print(battery_summary)
    
    return df

def main():
    """主函数：处理所有电池数据文件"""
    
    data_files = [
        'data_set/Dataset_NCA_battery.xlsx',
        'data_set/Dataset_NCM_battery.xlsx',
        'data_set/Dataset_NCM_NCA_battery.xlsx'
    ]
    
    for file_path in data_files:
        if os.path.exists(file_path):
            print(f"\n{'='*50}")
            print(f"处理文件: {file_path}")
            print(f"{'='*50}")
            try:
                add_initial_capacity_to_file(file_path)
            except Exception as e:
                print(f"处理文件 {file_path} 时出错: {e}")
        else:
            print(f"文件不存在: {file_path}")

if __name__ == "__main__":
    main()
